package strings;

import java.util.Queue;
import java.util.List;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.Set;
import java.util.HashSet;

/**
 * Dori & Landau version of Aho & Corasick automaton, implemented on
 * top of Bently & Sedgewick ternary search tree.
 * @author Dale Gerdemann
 */
public class AhoCorasickTST implements SetMatch {

    /** The callback. */
    private MatchAction ma;

    /** Size of TST. */
    private int size;

    /** Top of the trie. Added so that TST can represent empty string. */
    private Node superRoot;

    /** Root of TST. */
    private Node root;

    /** Characters for the top level transitions. Added to basic TST. */
    private Set<Character> topChars = new HashSet<Character>();

    /** Used as name for the nodes. Incremented each time a node is
     * constructed. */
    private static int index = 1;

    /** All nodes in the automaton. */
    private List<Node> allNode = new ArrayList<Node>();

    class Node {

        /** Constructor. */
        private Node() { }

        /** Next input character. */
        private char c;

        /** Left subtrie. */
        private Node left;
        /** Mid subtrie. */
        private Node mid;
        /** Right subtrie. */
        private Node right;

        /** The "value" associated with string (for debugging/pretty printing)
            where a string is identified with a node in the automaton, and the
            "value" is the name of that node. */
        private int val;

        /** Failure and output links as defined by Aho & Corasick. */
        private Node failure;

        /** Output for final states. */
        private int output = -1;
    }

    /**
     * Constructor.
     */
    public AhoCorasickTST() {
        superRoot = new Node();
        superRoot.val = 0;
        size = 0;      // initial size
        superRoot.failure = superRoot;
        allNode.add(superRoot);
    }


    /**
     * Compilation. To be performed after construction, and after
     * search words have been added.
     */
    public void compile() {
        addFailure();

    }

    /**
     * Search for patterns in input, depends on compile.
     * @param input the text to be searched
     */
    public void search(CharIterator input) {
        Node state = superRoot;

        int position = 0;

        while (input.hasNext()) {
            char c = input.next();
            position++;

            // feed c into automaton, go to the next state
            while (goTo(state, c) == null) {
                state = state.failure;
            }
            state = goTo(state, c);
            output(state, position);
        }
    }

    /**
     * Method to supply a callback.
     * @param ma0 the callback
     */
    public void setMatchAction(MatchAction ma0) {
        this.ma = ma0;
    }

    /**
     * The function described in articles by Dori & Landau for adding
     * failure links.
     */
    private void addFailure() {
        Queue<Node> q = new LinkedList<Node>();
        // Base case. Failure for depth 1 is the root node.
        for (char c : topChars) {
            Node x = goTo(superRoot, c);
            x.failure = superRoot;
            q.add(x);
        }
        // Traverse for nodes at depth greater than 1.
        while (!q.isEmpty()) {
            Node r = q.remove();
            traverse(r.mid, r, q);
        }
    }


    /**
     * Method described in Dori & Landau for adding failure links.
     * @param p current state in traversal
     * @param r the parent of p
     * @param q a container to store fail states
     */
    private void traverse(Node p, Node r, Queue<Node> q) {
        if (p == null) {
            return;
        }
        Node s = p;
        char a = s.c;
        Node state = r.failure;

        while (goTo(state, a) == null) {
            state = state.failure;
        }

        Node failstate = goTo(state, a);
        s.failure = failstate;
        q.add(s);
        traverse(s.left, r, q);
        traverse(s.right, r, q);
    }



    /**
     * The transition function described in articles by Aho & Corasick and
     * Dori & Landau. Modified to use the callback <code>ma</code>
     * (Match action).
     * @param x the current node, corresponds to a statte in the automaton
     * @param a the next input
     * @return destination node
     */
    private Node goTo(Node x, char a) {
        if (x == null) {
            return null;
        }
        if (x == superRoot) {
            if (topChars.contains(a)) {
                return get(root, "" + a, 0);
            } else {
                return superRoot;
            }
        }
        return get(x.mid, "" + a, 0);
    }


    /**
     * The output function described in articles by Aho & Corasick and
     * Dori & Landau. Modified to use the callback <code>ma</code>
     * (Match action).
     * @param state the state in the automaton
     * @param position the position in the input
     */
    private void output(Node state, int position) {
        while (true) {
            if (state == superRoot) {
                break;
            }
            if (state.output != -1) {
                ma.action(position, state.output);
            }
            state = state.failure;
        }
    }



    // Remainder of methods are specific to original TST.
    /**
     * @return number of key-value pairs
     */
    public int size() {
        return size;
    }


   /**************************************************************
    * Is string key in the symbol table?
    **************************************************************/

    /**
     * Access TST to find if value exists for given key.
     * @param key string to search for
     * @return true if key exists
     */
    public boolean contains(String key) {
        return get(key) != -1;
    }

    /**
     * Access TST to find value for given key.
     * @param key string to search for
     * @return int id for key
     */
    public int get(String key) {
        if (key == null || key.length() == 0) {
            throw new RuntimeException("illegal key");
        }
        Node x = get(root, key, 0);
        if (x == null) { return -1; }
        return x.output;
    }

    /**
     * Access TST to find value for given key.
     * @param x the Node in the TST
     * @param key string to search for
     * @param d index into key
     * @return subtrie corresponding to given key
     */
    private Node get(Node x, String key, int d) {
        if (key == null || key.length() == 0) {
            throw new RuntimeException("illegal key");
        }
        if (x == null) { return null; }
        char c = key.charAt(d);
        if      (c < x.c)              { return get(x.left,  key, d);
        } else if (c > x.c)            { return get(x.right, key, d);
        } else if (d < key.length() - 1) { return get(x.mid,   key, d + 1);
        } else                           { return x; }
    }


   /**************************************************************
    * Insert string s into the symbol table.
    * @param s string to insert
    **************************************************************/

    public void put(String s) {
        topChars.add(s.charAt(0));
        if (!contains(s)) {
            root = put(root, s, 0);
        }
    }

    /**
     * Add String to TST.
     * @param x the Node in the TST
     * @param s string to add
     * @param d index into s
     * @return updated Node
     */
    private Node put(Node x, String s, int d) {
        char c = s.charAt(d);
        if (x == null) {
            x = new Node();
            x.val = index++;   // Val used as id for referencing nodes
            x.c = c;
            allNode.add(x);
        }
        if      (c < x.c) {            x.left  = put(x.left,  s, d);
        } else if (c > x.c) {            x.right = put(x.right, s, d);
        } else if (d < s.length() - 1) { x.mid   = put(x.mid,   s, d + 1);
        } else {                         x.output   = size++; }
        return x;
    }

    /**
     * Add string (actualy CharIterator) to underlying trie. This
     * string then becomes a member of the set of strings to search.
     * @param cIter a sequence to search for.
     */
    public void put(CharIterator cIter) {
        StringBuilder sb = new StringBuilder();
        while (cIter.hasNext()) {
            sb.append(cIter.next());
        }
        put(sb.toString());
    }



   /**************************************************************
    * Find and return longest prefix of s in TST
    **************************************************************/

    /*
    public String longestPrefixOf(String s) {
        if (s == null || s.length() == 0) return null;
        int length = 0;
        Node x = root;
        int i = 0;
        while (x != null && i < s.length()) {
            char c = s.charAt(i);
            if      (c < x.c) x = x.left;
            else if (c > x.c) x = x.right;
            else {
                i++;
                if (x.val != null) length = i;
                x = x.mid;
            }
        }
        return s.substring(0, length);
    }
    */

    /**
     * Get iterator for keys.
     * @return an iterator over all keys (search words) in underlying Trie
     */
    public Iterable<String> keys() {
        Queue<String> queue = new LinkedList<String>();
        collect(root, "", queue);
        return queue;
    }

    /**
     * Find keys (search words) with a particuar prefix.
     * @param prefix required prefix of keys to return
     * @return all keys starting with given prefix
     */

    public Iterable<String> prefixMatch(String prefix) {
        Queue<String> queue = new LinkedList<String>();
        Node x = get(root, prefix, 0);
        if (x.val != -1) {
            queue.add(prefix);
        }
        collect(x.mid, prefix, queue);
        return queue;
    }

    /**
     * Collect all keys in subtrie rooted at x with given prefix.
     * @param x starting Node
     * @param prefix the required prefix
     * @param queue results collection
     */
    private void collect(Node x, String prefix, Queue<String> queue) {
        if (x == null) {
            return;
        }
        collect(x.left,  prefix,       queue);
        if (x.val != -1) {
            queue.add(prefix + x.c);
        }
        collect(x.mid,   prefix + x.c, queue);
        collect(x.right, prefix,       queue);
    }

    /**
     * Find keys (search words) matching a particuar wildcard pattern.
     * @param pat the pattern to match.
     * @return all keys matching given wildcard pattern.
     */
    public Iterable<String> wildcardMatch(String pat) {
        Queue<String> queue = new LinkedList<String>();
        collect(root, "", 0, pat, queue);
        return queue;
    }

    /**
     * Collector.
     * @param x tree node.
     * @param prefix the prefix to match.
     * @param i the input position.
     * @param pat todo.
     * @param q a queue to store results.
     */
    private void collect(Node x, String prefix, int i, String pat,
                        Queue<String> q) {
        if (x == null) {
            return;
        }
        char c = pat.charAt(i);
        if (c == '.' || c < x.c) {
            collect(x.left, prefix, i, pat, q);
        }
        if (c == '.' || c == x.c) {
            if (i == pat.length() - 1 && x.val != -1) {
                q.add(prefix + x.c);
            }
            if (i < pat.length() - 1) {
                collect(x.mid, prefix + x.c, i + 1, pat, q);
            }
        }
        if (c == '.' || c > x.c) {
            collect(x.right, prefix, i, pat, q);
        }
    }
}
